from tensorflow.keras import backend as K
from tensorflow.keras.layers import (Activation, BatchNormalization, Conv2D,
                                     DepthwiseConv2D, Dropout,
                                     GlobalAveragePooling2D, Input, Reshape)
from tensorflow.keras.models import Model


def _conv_block(inputs, filters, alpha, kernel=(3, 3), strides=(1, 1)):
    filters = int(filters * alpha)
    x = Conv2D(filters, kernel,
                      padding='same',
                      use_bias=False,
                      strides=strides,
                      name='conv1')(inputs)
    # BN层
    x = BatchNormalization(name='conv1_bn')(x)
    # 激活函数
    return Activation(relu6, name='conv1_relu')(x)

def _depthwise_conv_block(inputs, pointwise_conv_filters, alpha,
                          depth_multiplier=1, strides=(1, 1), block_id=1):

    pointwise_conv_filters = int(pointwise_conv_filters * alpha)

    # 首先对特征层进行深度可分离卷积
    x = DepthwiseConv2D((3, 3),
                        padding='same',
                        depth_multiplier=depth_multiplier,  # 每个输入通道的深度方向卷积输出通道的数量
                        strides=strides,  # 卷积的滑动步长。
                        use_bias=False,  # 不使用偏置向量
                        name='conv_dw_%d' % block_id)(inputs)
    # 标准化
    x = BatchNormalization(name='conv_dw_%d_bn' % block_id)(x)
    x = Activation(relu6, name='conv_dw_%d_relu' % block_id)(x)
    # 上面部分是深度可分离卷积
    # 下面是1x1的普通卷积，两个部分代替普通的3x3卷积
    # 进行1x1的普通卷积，对输入进来的特征层进行通道数的调整
    x = Conv2D(pointwise_conv_filters, (1, 1),
               padding='same',
               use_bias=False,
               strides=(1, 1),
               name='conv_pw_%d' % block_id)(x)
    # 再次进行标准化和激活函数
    x = BatchNormalization(name='conv_pw_%d_bn' % block_id)(x)
    return Activation(relu6, name='conv_pw_%d_relu' % block_id)(x)


def MobileNet(input_shape=None,
              alpha=0.25,
              depth_multiplier=1,
              dropout=1e-3,
              classes=1000):
    # 首先一张图片进入模型
    img_input = Input(shape=input_shape)

    # 224,224,3 -> 112,112,32
    # 首先进行卷积标准化加激活函数，通道数32 步长是2x2
    x = _conv_block(img_input, 32, alpha, strides=(2, 2))
    
    # 112,112,32 -> 112,112,64
    # 深度可分离卷积块
    x = _depthwise_conv_block(x, 64, alpha, depth_multiplier, block_id=1)
    # 因为步长是1x1，通道是64，所以特征层的高宽不会得到压缩

    # 112,112,64 -> 56,56,128
    # 再进行两次深度可分离卷积
    # 第一个步长是2x2，所以会压缩
    x = _depthwise_conv_block(x, 128, alpha, depth_multiplier,
                              strides=(2, 2), block_id=2)
    x = _depthwise_conv_block(x, 128, alpha, depth_multiplier, block_id=3)


    # 56,56,128 -> 28,28,256
    # 同上
    x = _depthwise_conv_block(x, 256, alpha, depth_multiplier,
                              strides=(2, 2), block_id=4)
    x = _depthwise_conv_block(x, 256, alpha, depth_multiplier, block_id=5)
    

    # 28,28,256 -> 14,14,512
    # 再进行6次深度可分离卷积块
    x = _depthwise_conv_block(x, 512, alpha, depth_multiplier,
                              strides=(2, 2), block_id=6)
    x = _depthwise_conv_block(x, 512, alpha, depth_multiplier, block_id=7)
    x = _depthwise_conv_block(x, 512, alpha, depth_multiplier, block_id=8)
    x = _depthwise_conv_block(x, 512, alpha, depth_multiplier, block_id=9)
    x = _depthwise_conv_block(x, 512, alpha, depth_multiplier, block_id=10)
    x = _depthwise_conv_block(x, 512, alpha, depth_multiplier, block_id=11)

    # 14,14,512 -> 7,7,1024
    # 再进行两次，同上
    x = _depthwise_conv_block(x, 1024, alpha, depth_multiplier,
                              strides=(2, 2), block_id=12)
    x = _depthwise_conv_block(x, 1024, alpha, depth_multiplier, block_id=13)

    # 7,7,1024 -> 1,1,1024
    # 进行一个全局平均池化，得到一个长度为1024的特征长条
    x = GlobalAveragePooling2D()(x)

    # 转化为特征层
    shape = (1, 1, int(1024 * alpha))

    x = Reshape(shape, name='reshape_1')(x)
    # 通过忽略部分特征检测器来减少隐层节点间的作用，防止过拟合
    x = Dropout(dropout, name='dropout')(x)
    # 通过卷积调整通道数，相当于对输入进来的图片进行分类
    x = Conv2D(classes, (1, 1), padding='same', name='conv_preds')(x)
    # 获得每一个种类的概率
    x = Activation('softmax', name='act_softmax')(x)
    x = Reshape((classes,), name='reshape_2')(x)

    inputs = img_input

    model = Model(inputs, x, name='mobilenet_%0.2f' % (alpha))
    return model

# 用作神经元的输出，修正线性单元
def relu6(x):
    return K.relu(x, max_value=6)

# 显示输出网络结构
if __name__ == '__main__':
    model = MobileNet(input_shape=(224, 224, 3),classes=2,alpha=0.25)
    model.summary()
